package register;

import mathtools.Complex;
import mathtools.Matrix;
import exceptions.MatrixSizeException;

/**
 * Class Register is a one-dimensional Matrix which will hold the amplitudes of the Quantum 
 * state. 
 * @author Richard Inskip, Matt Cole
 *
 */

public class Register extends Matrix {

	/** Number of states in the register */
	private int numStates;
	/** Number of qubits in the register */
	private int numQbits;

	/**
	 * Constructor to create the Quantum Register depending on the number of Quantum bits.
	 * @param numQubits The number of Quantum bits.
	 */
	public Register(int numQubits){
		super(1, (int)Math.pow(2,numQubits), 0);
		numStates = (int)Math.pow(2,numQubits);
		this.numQbits = numQubits;
		setToGroundState();
	}
	
	/**
	 * Puts the register into the ground state. Called when the register is created and also used
	 * when the reset button is clicked on the UI.
	 */
	public void setToGroundState() {
		m[0][0] = new Complex(1,0);
		for	(int j = 1; j < numStates; j++){
			m[0][j] = new Complex(0,0);
		}
	}

	/**
	 *	 * A Method to update the Register with a new version. This will allow the Register object
	 * to always update itself.
	 * @param r A new 1-dimensional Vector to update the elements in the Register.
	 * @throws MatrixSizeException 
	 */
	public void update(Matrix r){
		/* Make sure the updating vector is of the same dimensions as the
		 * current register.
		 */
		if(r.getHeight()!= height && r.getWidth()!= width){
			throw new IllegalArgumentException("The updating Register must be of exact size as the current Register.\n" +
					"New register size: " + r.getHeight() + "Old register size: " + height);
		}
		else{
			// Set the elements of the updating register into the register.
			for(int j = 0; j < numStates; j++){
				setElement(0,j,r.getElement(0, j));
			}
		}
	}

	/**
	 * A Method to update the Register.
	 * @param r A new 1-dimensional Vector to update the elements in the Register.
	 */
	public void update(Complex [] r){
		/* Make sure the updating vector is of the same dimensions as the
		 * current register. If the dimensions are different, throw a 
		 * exception (TODO).  
		 */
		if(r.length != height && r.length != width){
			throw new IllegalArgumentException("The updating Register must be of exact size as the current Register.\n" +
					"New register size: " + r.length + "Old register size: " + getHeight());
		}

		// Set the elements of the updating register into the register.
		for(int j = 0; j < numStates; j++){
			setElement(0, j, r[ j]);
		}
	}

	/** 
	 * A Method to normalise the amplitudes of the Register.
	 * @throws MatrixSizeException 
	 */
	public void normalise(){
		double sum = 0.0;
		for(int j = 0; j < numStates; j++){
			sum += getElement(0,j).getMagnitudeSquared();

		}
		float normalisation_constant = (float)(1/Math.sqrt(sum));
		Matrix normalised_register = multiplyMatrix(normalisation_constant);
		update(normalised_register);
	}
	/**
	 * Return the number of qubits in the register. This is NOT the number of states (or amplitudes)
	 * @return numQubits Number of qubits represented by the register
	 */
	public int getNumQbits() {
		return numQbits;
	}

	/**
	 * Return the number of states in the register, or number of amplitudes
	 * @return numberStates Number of states in the register
	 */
	public int getNumStates(){
		return numStates;
	}
	
	/**
	 * Return the probability of a given state. Needed for drawing graph of probabilities
	 * @param state State to give probability for
	 * @return probability Probability of measuring given state
	 */
	public double getProbabilityOfState(int state){
		return m[0][state].getMagnitudeSquared();
	}
	
	
	/**
	 * Post measurement. Returns the state that has probability 1. For use only after the measurement gate has been applied to the register.
	 * @return i State that has probability 1 or hence measure state
	 */
	public int postMeasure() {
		for(int i = 0; i < numStates; i++) {
			if (this.getProbabilityOfState(i)==1 ){
				return i;
			}
		}
		throw new RuntimeException("Register has not been measured correctly. No certain state.");
	}
}


